from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.agents.llm_agent import LlmAgent
from google.genai import types
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner

from google.adk.models.lite_llm import LiteLlm

# --- Constants ---
APP_NAME = "code_pipeline_app"
USER_ID = "dev_user_01"
SESSION_ID = "pipeline_session_01"
GEMINI_MODEL = "deepseek/deepseek-chat"

# --- 1. Define Sub-Agents for Each Pipeline Stage ---

# Code Writer Agent
# Takes the initial specification (from user query) and writes code.
code_writer_agent = LlmAgent(
    name="CodeWriterAgent",
    model=LiteLlm(model="deepseek/deepseek-chat"),
    instruction="""You are a Code Writer AI.
    Based on the user's request, write the initial Python code.
    Output *only* the raw code block. 最终的结果使用中文回答
    """,
    description="Writes initial code based on a specification.",
    # Stores its output (the generated code) into the session state
    # under the key 'generated_code'.
    output_key="generated_code"
)

# Code Reviewer Agent
# Takes the code generated by the previous agent (read from state) and provides feedback.
code_reviewer_agent = LlmAgent(
    name="CodeReviewerAgent",
    model=LiteLlm(model="deepseek/deepseek-chat"),
    instruction="""You are a Code Reviewer AI.
    Review the Python code provided in the session state under the key 'generated_code'.
    Provide constructive feedback on potential errors, style issues, or improvements.
    Focus on clarity and correctness.
    Output only the review comments. 最终的结果使用中文回答
    """,
    description="Reviews code and provides feedback.",
    # Stores its output (the review comments) into the session state
    # under the key 'review_comments'.
    output_key="review_comments"
)

# Code Refactorer Agent
# Takes the original code and the review comments (read from state) and refactors the code.
code_refactorer_agent = LlmAgent(
    name="CodeRefactorerAgent",
    model=LiteLlm(model="deepseek/deepseek-chat"),
    instruction="""You are a Code Refactorer AI.
    Take the original Python code provided in the session state key 'generated_code'
    and the review comments found in the session state key 'review_comments'.
    Refactor the original code to address the feedback and improve its quality.
    Output *only* the final, refactored code block. 最终的结果使用中文回答
    """,
    description="Refactors code based on review comments.",
    # Stores its output (the refactored code) into the session state
    # under the key 'refactored_code'.
    output_key="refactored_code"
)

# --- 2. Create the SequentialAgent ---
# This agent orchestrates the pipeline by running the sub_agents in order.
code_pipeline_agent = SequentialAgent(
    name="CodePipelineAgent",
    sub_agents=[code_writer_agent, code_reviewer_agent, code_refactorer_agent]
    # The agents will run in the order provided: Writer -> Reviewer -> Refactorer
)

# Session and Runner
session_service = InMemorySessionService()
session = session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID)
runner = Runner(agent=code_pipeline_agent, app_name=APP_NAME, session_service=session_service)

# Export the agent for use in other modules
agent = code_pipeline_agent